# 导入必要的库
import os  # 操作系统接口，用于文件路径操作
import torch  # PyTorch深度学习框架
import torch.nn as nn  # 神经网络模块
import torch.optim as optim  # 优化算法模块
import torchvision  # 计算机视觉相关工具
import torchvision.transforms as transforms  # 图像预处理工具
import matplotlib.pyplot as plt  # 数据可视化库
from torch.optim.lr_scheduler import CosineAnnealingLR  # 余弦退火学习率调度器

# ============================ 1. 训练超参数 ============================
BATCH_SIZE = 128  # 每个训练批次的样本数量
NUM_EPOCHS = 100  # 总训练轮数
LEARNING_RATE = 0.1  # 初始学习率
MODEL_SAVE_PATH = './results/cifar100_model.pth'  # 模型保存路径

# CIFAR-100 数据集的标准化参数（RGB三个通道的均值和标准差）
MEAN = (0.5071, 0.4865, 0.4409)  # 均值
STD = (0.2673, 0.2564, 0.2762)  # 标准差

# ============================ 2. 数据增强（防止过拟合） ============================
# 训练集数据增强组合
transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),  # 随机裁剪（32x32）并填充4像素边界
    transforms.RandomHorizontalFlip(),  # 以50%概率水平翻转
    transforms.RandomRotation(15),  # 在±15度范围内随机旋转
    transforms.ColorJitter(brightness=0.2,  # 颜色抖动：亮度调整幅度
                           contrast=0.2,  # 对比度调整幅度
                           saturation=0.2,  # 饱和度调整幅度
                           hue=0.1),  # 色相调整幅度（范围-0.1到0.1）
    transforms.ToTensor(),  # 将PIL图像转换为Tensor格式（值范围0-1），必须放在Normalize之前
    transforms.Normalize(MEAN, STD),  # 标准化处理（减去均值，除以标准差）
    transforms.RandomErasing(p=0.2, scale=(0.02, 0.2))  # 随机擦除（概率20%，擦除面积比例2%-20%）
])

# 测试集数据预处理（不需要数据增强）
transform_test = transforms.Compose([
    transforms.ToTensor(),  # 转换为Tensor
    transforms.Normalize(MEAN, STD)  # 使用相同的标准化参数
])

# 下载并加载CIFAR-100数据集
trainset = torchvision.datasets.CIFAR100(
    root='./data',  # 数据集存储路径
    train=True,  # 加载训练集
    download=True,  # 如果本地不存在则自动下载
    transform=transform_train  # 应用训练集的数据增强
)
# 创建训练数据加载器
trainloader = torch.utils.data.DataLoader(
    trainset,
    batch_size=BATCH_SIZE,
    shuffle=True,  # 打乱训练数据顺序
    num_workers=0  # 数据加载线程数（0表示使用主线程）
)

# 加载测试集
testset = torchvision.datasets.CIFAR100(
    root='./data',
    train=False,  # 加载测试集
    download=True,
    transform=transform_test  # 应用测试集的预处理
)
# 创建测试数据加载器
testloader = torch.utils.data.DataLoader(
    testset,
    batch_size=BATCH_SIZE,
    shuffle=False,  # 测试集不需要打乱顺序
    num_workers=0
)


# ============================ 3. ResNet 结构 ============================
class ResidualBlock(nn.Module):
    """ResNet残差块，实现恒等映射（identity mapping）"""

    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        """
        参数说明:
        in_channels: 输入通道数
        out_channels: 输出通道数
        stride: 卷积步长（默认为1）
        downsample: 下采样方法（当维度不匹配时需要）
        """
        super(ResidualBlock, self).__init__()
        # 第一个卷积层：3x3卷积，可能改变特征图尺寸
        self.conv1 = nn.Conv2d(
            in_channels, out_channels,
            kernel_size=3, stride=stride,
            padding=1, bias=False  # 不使用偏置（因为后面接BN层）
        )
        self.bn1 = nn.BatchNorm2d(out_channels)  # 批标准化
        # 第二个卷积层：3x3卷积，保持特征图尺寸
        self.conv2 = nn.Conv2d(
            out_channels, out_channels,
            kernel_size=3, padding=1, bias=False
        )
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)  # 原地激活（节省内存）
        self.downsample = downsample  # 下采样方法（当输入输出维度不一致时需要）

    def forward(self, x):
        identity = x  # 保留原始输入作为残差连接

        # 如果存在下采样（如通道数变化或尺寸变化）
        if self.downsample is not None:
            identity = self.downsample(x)  # 对原始输入进行维度调整

        # 第一个卷积块
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        # 第二个卷积块
        out = self.conv2(out)
        out = self.bn2(out)

        # 残差连接
        out += identity  # 核心，解决梯度消失问题：将调整后的原始输入与卷积结果相加
        out = self.relu(out)  # 最后再激活一次

        return out


class ResNet(nn.Module):
    """改进的ResNet-18结构（适用于CIFAR-100的32x32输入尺寸）"""

    def __init__(self, num_classes=100):
        super(ResNet, self).__init__()
        # 初始卷积层（原ResNet使用7x7卷积+池化，这里改为3x3卷积保持分辨率）
        self.conv1 = nn.Conv2d(
            3, 64,  # 输入通道3（RGB），输出通道64
            kernel_size=3, padding=1, bias=False
        )
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)

        # 构建残差层（三层残差层，每层两个残差块）
        self.layer1 = self._make_layer(64, 128, 2)  # 第1层：2个残差块
        self.layer2 = self._make_layer(128, 256, 2, stride=2)  # 第2层：2个残差块，下采样
        self.layer3 = self._make_layer(256, 512, 2, stride=2)  # 第3层：2个残差块，下采样

        # 全局平均池化（将特征图尺寸降为1x1）
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        # 全连接层（512维特征 → 100类）
        self.fc = nn.Linear(512, num_classes)

    def _make_layer(self, in_channels, out_channels, blocks, stride=1):
        """构建包含多个残差块的层"""
        downsample = None
        # 当输入输出维度不一致时（通道数变化或特征图尺寸变化），需要下采样
        if stride != 1 or in_channels != out_channels:
            downsample = nn.Sequential(
                # 1x1卷积调整通道数和尺寸
                nn.Conv2d(
                    in_channels, out_channels,
                    kernel_size=1, stride=stride, bias=False
                ),
                nn.BatchNorm2d(out_channels)
            )

        layers = []
        # 第一个残差块可能需要下采样
        layers.append(ResidualBlock(in_channels, out_channels, stride, downsample))
        # 后续残差块（保持输入输出维度一致）
        for _ in range(1, blocks):
            layers.append(ResidualBlock(out_channels, out_channels))

        return nn.Sequential(*layers)  # 将层列表转换为顺序容器

    def forward(self, x):
        # 初始卷积
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)

        # 通过各个残差层
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)

        # 全局平均池化
        x = self.avgpool(x)
        # 展平特征（保留批量维度）
        x = torch.flatten(x, 1)
        # 全连接层分类
        x = self.fc(x)
        return x


# ============================ 4. 训练设置 ============================
# 选择设备（优先使用GPU）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 实例化模型并转移到设备
model = ResNet().to(device)

# 损失函数：交叉熵损失（带标签平滑，系数0.1）
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
# 优化器：带动量的SGD
optimizer = optim.SGD(
    model.parameters(),
    lr=LEARNING_RATE,
    momentum=0.9,  # 动量系数
    weight_decay=5e-4  # L2正则化系数
)
# 学习率调度器：余弦退火（周期为总训练轮数）
scheduler = CosineAnnealingLR(optimizer, T_max=NUM_EPOCHS)

# ============================ 5. 训练和测试 ============================
# 用于记录训练过程的列表
loss_list = []  # 每个epoch的训练损失
train_acc_list = []  # 每个epoch的训练准确率
test_acc_list = []  # 每个epoch的测试准确率


def train_one_epoch():
    """训练单个epoch"""
    model.train()  # 设置为训练模式（启用BatchNorm和Dropout）
    correct = 0  # 正确预测的样本数
    total = 0  # 总样本数
    running_loss = 0.0  # 累计损失

    for inputs, labels in trainloader:
        # 将数据转移到设备（GPU/CPU）
        inputs, labels = inputs.to(device), labels.to(device)

        # 梯度清零
        optimizer.zero_grad()

        # 前向传播
        outputs = model(inputs)
        # 计算损失
        loss = criterion(outputs, labels)

        # 反向传播
        loss.backward()
        # 参数更新
        optimizer.step()

        # 统计信息
        running_loss += loss.item() * inputs.size(0)  # 累加损失（考虑批次大小）
        _, predicted = outputs.max(1)  # 获取预测类别
        total += labels.size(0)  # 累加样本总数
        correct += predicted.eq(labels).sum().item()  # 累加正确数

    # 计算平均损失和准确率
    epoch_loss = running_loss / total
    epoch_acc = correct / total
    return epoch_loss, epoch_acc


def evaluate():
    """在测试集上评估模型性能"""
    model.eval()  # 设置为评估模式（关闭BatchNorm和Dropout）
    correct = 0
    total = 0

    with torch.no_grad():  # 禁用梯度计算（节省内存）
        for inputs, labels in testloader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()

    return correct / total  # 返回测试准确率


# ============================ 6. 训练循环 ============================
for epoch in range(NUM_EPOCHS):
    print(f"Epoch {epoch + 1}/{NUM_EPOCHS}:")

    # 训练一个epoch并获取损失和准确率
    train_loss, train_acc = train_one_epoch()
    print(f"    Train Loss: {train_loss:.4f}, Train Accuracy: {train_acc:.4f}")

    # 在测试集上评估
    test_acc = evaluate()
    print(f"    Test Accuracy: {test_acc:.4f}")

    # 更新学习率（根据余弦退火策略）
    scheduler.step()

    # 记录训练过程数据
    loss_list.append(train_loss)
    train_acc_list.append(train_acc)
    test_acc_list.append(test_acc)

# ============================ 7. 保存模型 ============================
# 创建保存目录（如果不存在）
os.makedirs(os.path.dirname(MODEL_SAVE_PATH), exist_ok=True)
# 保存模型权重
torch.save(model.state_dict(), MODEL_SAVE_PATH)
print(f"模型已保存至 {MODEL_SAVE_PATH}")

# ============================ 8. 保存训练过程曲线 ============================
epochs = range(1, NUM_EPOCHS + 1)  # x轴数据（1到总轮数）

# 绘制准确率曲线
plt.figure()
plt.plot(epochs, train_acc_list, label='Train Accuracy')
plt.plot(epochs, test_acc_list, label='Test Accuracy')
plt.xlabel('Epoch')  # x轴标签
plt.ylabel('Accuracy')  # y轴标签
plt.title('Train and Test Accuracy')  # 标题
plt.legend()  # 显示图例
plt.savefig("accuracy_curve.png")  # 保存为图片
plt.close()  # 关闭当前图像，防止内存泄漏

# 绘制损失曲线
plt.figure()
plt.plot(epochs, loss_list, label='Train Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training Loss')
plt.legend()
plt.savefig("loss_curve.png")
plt.close()

print("训练曲线图片已保存至 accuracy_curve.png 和 loss_curve.png")
